// SPDX-License-Identifier: BSD-2-Clause
/*
 * Copyright (c) 2018, Linaro Limited
 */

#include <stdio.h>
#include <stdlib.h>
#include "talib1.h"
#include "kem.h"
#include "sign.h"
#include "keyfind.h"
#include "mtwister.h"
#include <bn_pairing.h>
#include <tee_internal_api.h>
#include <libs_ta.h>
#include "utility.h"
#include "string.h"
#define TA_SECURE_STORAGE_UUID \
		{ 0xf4e750bb, 0x1437, 0x4fbf, \
			{ 0x87, 0x85, 0x8d, 0x35, 0x80, 0xc3, 0x49, 0x94 } }

#define TA_SECURE_STORAGE_CMD_READ_RAW		0
#define TA_SECURE_STORAGE_CMD_WRITE_RAW		1
#define TA_SECURE_STORAGE_CMD_DELETE		2

#define MLEN 59
#define NTESTS 10000

uint8_t Kem_Pk[CRYPTO_PUBLICKEYBYTES];
uint8_t Kem_Sk[CRYPTO_SECRETKEYBYTES];
uint8_t Self_Sig_Pk[CRYPTO_PUBLICKEYBYTES_SIG];
uint8_t Self_Sig_Sk[CRYPTO_SECRETKEYBYTES_SIG];
uint8_t Qrng_Random[650];
TEE_Result TA_CreateEntryPoint(void)
{
	printf("TA_CreateEntryPoint()\n");
	return TEE_SUCCESS;
}

void TA_DestroyEntryPoint(void)
{
	printf("TA_DestroyEntryPoint()\n");
}

TEE_Result TA_OpenSessionEntryPoint(uint32_t __unused param_types,
		TEE_Param __unused params[4],
		void __unused **sess_ctx)
{
	printf("TA_OpenSession()\n");
	return TEE_SUCCESS;
}

void TA_CloseSessionEntryPoint(void __unused *sess_ctx)
{
	printf("TA_CloseSession()\n");
}


void secure_write()
{
	TEE_Result res1;
	TEE_TASessionHandle sess;
	TEE_UUID uuid = TA_SECURE_STORAGE_UUID;
	uint32_t time=10;
	uint32_t ret_origin;
	uint32_t param_types1;
	TEE_Param params1[4];
	//// TCP datapackage
	uint8_t package[32];
	uint8_t package_id[8];
	TEE_MemMove(package_id, node1.id, 8);
	TEE_MemMove(package, node1.id, 8);
	TEE_MemMove(package+8, node1.rid, 8);
	TEE_MemMove(package+16, node1.sk, 16);
	printf("\nSK storage start:\n");
	printf("storage_id:\n");
	format_print(package_id, 8);
	printf("storage_package:\n");
	format_print(package, 32);
	// printf("\npackage:\n");
	// format_print(package, 32);
	param_types1 = TEE_PARAM_TYPES(TEE_PARAM_TYPE_MEMREF_INPUT,
				TEE_PARAM_TYPE_MEMREF_INPUT,
				TEE_PARAM_TYPE_NONE,
				TEE_PARAM_TYPE_NONE);
	params1[0].memref.buffer = package_id;
	params1[0].memref.size = 8;
	params1[1].memref.buffer = package;
	params1[1].memref.size = 32;
	res1 = TEE_OpenTASession(&uuid, time, param_types1, params1, &sess, &ret_origin);
	if (res1 != TEE_SUCCESS)
		printf("TEE_Opensession failed with code 0x%x origin 0x%x",
			res1, ret_origin);
	//printf("\nwrong6\n");
	res1 = TEE_InvokeTACommand(sess,time, TA_SECURE_STORAGE_CMD_WRITE_RAW, param_types1,params1,&ret_origin);
	if (res1 != TEE_SUCCESS)
		printf( "TEE_Invokecommand failed with code 0x%x origin 0x%x",
			res1, ret_origin);
	// for(int i=0;i<16;i++){
	// 	printf("%02x  ",read[i]);
	// }
	printf("\nSk storage successful!\n");
	TEE_CloseTASession(sess);
}

void secure_delete(uint8_t* delete_id)
{
	TEE_Result res1;
	TEE_TASessionHandle sess;
	TEE_UUID uuid = TA_SECURE_STORAGE_UUID;
	uint32_t time=10;
	uint32_t ret_origin;
	uint32_t param_types1;
	TEE_Param params1[4];
	param_types1 = TEE_PARAM_TYPES(TEE_PARAM_TYPE_MEMREF_INPUT,
				TEE_PARAM_TYPE_NONE,
				TEE_PARAM_TYPE_NONE,
				TEE_PARAM_TYPE_NONE);
	params1[0].memref.buffer = delete_id;
	params1[0].memref.size = 8;
	res1 = TEE_OpenTASession(&uuid, time, param_types1, params1, &sess, &ret_origin);
	if (res1 != TEE_SUCCESS)
		printf("TEE_Opensession failed with code 0x%x origin 0x%x",
			res1, ret_origin);
	//printf("\nwrong6\n");
	res1 = TEE_InvokeTACommand(sess,time, TA_SECURE_STORAGE_CMD_DELETE, param_types1,params1,&ret_origin);
	if (res1 != TEE_SUCCESS)
		printf( "TEE_Invokecommand failed with code 0x%x origin 0x%x",
			res1, ret_origin);
}

TEE_Result ta_save_opposigpk(uint32_t param_types, TEE_Param params[4])
{
	const uint32_t exp_param_types =
		TEE_PARAM_TYPES(TEE_PARAM_TYPE_MEMREF_INPUT,
			TEE_PARAM_TYPE_MEMREF_INPUT,
			TEE_PARAM_TYPE_NONE,
			TEE_PARAM_TYPE_NONE
			);
	TEE_Result res;
	if (param_types != exp_param_types)
		return TEE_ERROR_BAD_PARAMETERS;

	uint8_t oppo_id[ID_LEN];
	uint8_t oppo_sig_pk[CRYPTO_PUBLICKEYBYTES_SIG];
	TEE_MemMove(oppo_id, params[0].memref.buffer, ID_LEN);
	TEE_MemMove(oppo_sig_pk, params[1].memref.buffer, CRYPTO_PUBLICKEYBYTES_SIG);
	
	//save
	sig_data_insert(&sig_node_head, oppo_id, oppo_sig_pk);
}

TEE_Result ta_sig_keypair(uint32_t param_types, TEE_Param params[4])
{
	const uint32_t exp_param_types =
		TEE_PARAM_TYPES(TEE_PARAM_TYPE_MEMREF_OUTPUT,
			TEE_PARAM_TYPE_NONE,
			TEE_PARAM_TYPE_NONE,
			TEE_PARAM_TYPE_NONE
			);
	TEE_Result res;
	if (param_types != exp_param_types)
		return TEE_ERROR_BAD_PARAMETERS;
	uint8_t self_sig_pk[CRYPTO_PUBLICKEYBYTES_SIG];
  	uint8_t self_sig_sk[CRYPTO_SECRETKEYBYTES_SIG];
	crypto_sign_keypair(self_sig_pk, self_sig_sk);
	TEE_MemMove(params[0].memref.buffer, self_sig_pk, CRYPTO_PUBLICKEYBYTES_SIG);
	TEE_MemMove(Self_Sig_Pk, self_sig_pk, CRYPTO_PUBLICKEYBYTES_SIG);
	TEE_MemMove(Self_Sig_Sk, self_sig_sk, CRYPTO_SECRETKEYBYTES_SIG);
	return TEE_SUCCESS;
}

TEE_Result ta_kem_keypair(uint32_t param_types, TEE_Param params[4])
{
	const uint32_t exp_param_types =
		TEE_PARAM_TYPES(TEE_PARAM_TYPE_MEMREF_OUTPUT,
			TEE_PARAM_TYPE_NONE,
			TEE_PARAM_TYPE_NONE,
			TEE_PARAM_TYPE_NONE
			);
	TEE_Result res;
	if (param_types != exp_param_types)
		return TEE_ERROR_BAD_PARAMETERS;
	//kem_keypair
	uint8_t kem_pk[CRYPTO_PUBLICKEYBYTES];
	uint8_t kem_sk[CRYPTO_SECRETKEYBYTES];
	int ret = crypto_kem_keypair(kem_pk, kem_sk, Qrng_Random);
	TEE_MemMove(Kem_Pk, kem_pk, CRYPTO_PUBLICKEYBYTES);
	TEE_MemMove(Kem_Sk, kem_sk, CRYPTO_SECRETKEYBYTES);
	//sig_kem_pk
	uint8_t sig_msg[CRYPTO_PUBLICKEYBYTES + CRYPTO_BYTES_SIG];
	size_t sig_msg_len;
	uint8_t self_sig_sk[CRYPTO_SECRETKEYBYTES_SIG];
	TEE_MemMove(self_sig_sk, Self_Sig_Sk, CRYPTO_SECRETKEYBYTES_SIG);
	crypto_sign(sig_msg, &sig_msg_len, kem_pk, CRYPTO_PUBLICKEYBYTES, self_sig_sk);
	//return the result
	TEE_MemMove(params[0].memref.buffer, sig_msg, CRYPTO_PUBLICKEYBYTES + CRYPTO_BYTES_SIG);
	return TEE_SUCCESS;
}

TEE_Result ta_authen_encape(uint32_t param_types, TEE_Param params[4])
{
	const uint32_t exp_param_types =
		TEE_PARAM_TYPES(TEE_PARAM_TYPE_MEMREF_INPUT,
			TEE_PARAM_TYPE_MEMREF_INPUT,
			TEE_PARAM_TYPE_MEMREF_OUTPUT,
			TEE_PARAM_TYPE_NONE
			);
	TEE_Result res;
	if (param_types != exp_param_types)
		return TEE_ERROR_BAD_PARAMETERS;
	
	//sign_verify
	uint8_t oppo_sig_pk[CRYPTO_PUBLICKEYBYTES_SIG];
	uint8_t oppo_id[ID_LEN];
	uint8_t kem_pk[CRYPTO_PUBLICKEYBYTES];
	size_t kem_pk_len;
	uint8_t *sig_msg;
	int ret;
	size_t sig_msg_len = CRYPTO_PUBLICKEYBYTES + CRYPTO_BYTES_SIG;
	sig_msg = TEE_Malloc(sig_msg_len, 0);
	if(!sig_msg)
		return TEE_ERROR_OUT_OF_MEMORY;
	TEE_MemMove(sig_msg, params[0].memref.buffer, sig_msg_len);
	TEE_MemMove(oppo_id, params[1].memref.buffer, ID_LEN);
	ret = sig_data_search(&sig_node_head, oppo_id, oppo_sig_pk);
	if(ret) {
      printf("Found failed\n");
	  return -1;
    }
	ret = crypto_sign_open(kem_pk, &kem_pk_len, sig_msg, sig_msg_len, oppo_sig_pk);
	if(ret) {
      printf("Verification failed\n");
	  return -1;
    }

	//encape
	uint8_t ct[CRYPTO_CIPHERTEXTBYTES];
	uint8_t final_key[CRYPTO_BYTES];
	ret = crypto_kem_enc(ct, final_key, kem_pk);
	printf("\nfinal_key:\n");
	format_print(final_key, 16);

	//sign_ct(ciphertext)
	uint8_t *self_sig_sk;
	uint8_t *sig_ct;
	size_t sig_ct_len;
	self_sig_sk = TEE_Malloc(CRYPTO_SECRETKEYBYTES_SIG, 0);
	sig_ct = TEE_Malloc(CRYPTO_CIPHERTEXTBYTES + CRYPTO_BYTES_SIG, 0);
	if(!self_sig_sk || !sig_ct){
		return TEE_ERROR_OUT_OF_MEMORY;
	}
	TEE_MemMove(self_sig_sk, Self_Sig_Sk, CRYPTO_SECRETKEYBYTES_SIG);
	crypto_sign(sig_ct, &sig_ct_len, ct, CRYPTO_CIPHERTEXTBYTES, self_sig_sk);
	
	//return
	TEE_MemMove(node1.id, oppo_id, 8);
	// TEE_MemMove(node1.rid, sig_ct+1000, 8);	
	TEE_MemMove(node1.sk, final_key, 16);
	secure_write();
	TEE_MemMove(params[2].memref.buffer, sig_ct, sig_ct_len);
	TEE_Free(sig_msg);
	TEE_Free(sig_ct);
	TEE_Free(self_sig_sk);
}

TEE_Result ta_authen_decape(uint32_t param_types, TEE_Param params[4])
{
	const uint32_t exp_param_types =
		TEE_PARAM_TYPES(TEE_PARAM_TYPE_MEMREF_INPUT,
			TEE_PARAM_TYPE_MEMREF_INPUT,
			TEE_PARAM_TYPE_MEMREF_OUTPUT,
			TEE_PARAM_TYPE_NONE
			);
	TEE_Result res;
	if (param_types != exp_param_types)
		return TEE_ERROR_BAD_PARAMETERS;

	//sign_verify
	int ret;
	size_t ct_len;
	size_t sig_ct_len = CRYPTO_CIPHERTEXTBYTES + CRYPTO_BYTES_SIG;
	uint8_t ct[CRYPTO_CIPHERTEXTBYTES + CRYPTO_BYTES_SIG];
	uint8_t oppo_id[ID_LEN];
	uint8_t *oppo_sig_pk;
	uint8_t *sig_ct;
	uint8_t auth_res[] = "00";
	// uint8_t *auth_res;
	oppo_sig_pk = TEE_Malloc(CRYPTO_PUBLICKEYBYTES_SIG, 0);
	sig_ct = TEE_Malloc(CRYPTO_CIPHERTEXTBYTES + CRYPTO_BYTES_SIG, 0);
	// auth_res = TEE_Malloc(1, 0);
	if(!oppo_sig_pk || !sig_ct){
		return TEE_ERROR_OUT_OF_MEMORY;
	}
	TEE_MemMove(sig_ct, params[0].memref.buffer, CRYPTO_CIPHERTEXTBYTES + CRYPTO_BYTES_SIG);
	TEE_MemMove(oppo_id, params[1].memref.buffer, ID_LEN);

	ret = sig_data_search(&sig_node_head, oppo_id, oppo_sig_pk);
	if(ret) {
      printf("Found failed\n");
	  return -1;
    }
	ret = crypto_sign_open(ct, &ct_len, sig_ct, sig_ct_len, oppo_sig_pk);
	if(ret) {
      printf("Verification failed\n");
	  TEE_MemMove(params[2].memref.buffer, auth_res, 1);
	  return -1;
    }
	auth_res[0] = '1';
	TEE_MemMove(params[2].memref.buffer, auth_res, 1);

	//decape
	uint8_t kem_sk[CRYPTO_SECRETKEYBYTES];
	uint8_t final_key[CRYPTO_BYTES];
	TEE_MemMove(kem_sk, Kem_Sk, CRYPTO_SECRETKEYBYTES);

	ret = crypto_kem_dec(final_key, ct, kem_sk);
	printf("\nfinal_key:\n");
	format_print(final_key, 16);
	
	//return
	TEE_MemMove(node1.id, oppo_id, 8);
	TEE_MemMove(node1.rid, ct+2000, 8);
	TEE_MemMove(node1.sk, final_key, 16);
	secure_write();
	TEE_Free(oppo_sig_pk);
	TEE_Free(sig_ct);
	// TEE_Free(auth_res);
}

TEE_Result ta_sign(uint32_t param_types, TEE_Param params[4])
{
	const uint32_t exp_param_types =
		TEE_PARAM_TYPES(TEE_PARAM_TYPE_MEMREF_INPUT,
			TEE_PARAM_TYPE_MEMREF_INPUT,
			TEE_PARAM_TYPE_MEMREF_OUTPUT,
			TEE_PARAM_TYPE_MEMREF_OUTPUT
			);
	TEE_Result res;
	if (param_types != exp_param_types)
		return TEE_ERROR_BAD_PARAMETERS;
	
	uint8_t oppo_sig_pk[CRYPTO_PUBLICKEYBYTES_SIG];
	uint8_t oppo_id[ID_LEN];
	uint8_t *sig_msg;
	int ret;
	size_t sig_msg_len = 2;
	sig_msg = TEE_Malloc(sig_msg_len, 0);
	if(!sig_msg)
		return TEE_ERROR_OUT_OF_MEMORY;
	TEE_MemMove(sig_msg, params[0].memref.buffer, sig_msg_len);
	TEE_MemMove(oppo_id, params[1].memref.buffer, ID_LEN);

	//sign_ct(ciphertext)
	uint8_t *self_sig_sk;
	uint8_t *sig_ct;
	size_t sig_ct_len = sig_msg_len + CRYPTO_BYTES_SIG;
	self_sig_sk = TEE_Malloc(CRYPTO_SECRETKEYBYTES_SIG, 0);
	sig_ct = TEE_Malloc(sig_msg_len + CRYPTO_BYTES_SIG, 0);
	if(!self_sig_sk || !sig_ct){
		return TEE_ERROR_OUT_OF_MEMORY;
	}
	TEE_MemMove(self_sig_sk, Self_Sig_Sk, CRYPTO_SECRETKEYBYTES_SIG);
	crypto_sign(sig_ct, &sig_ct_len, sig_msg, sig_msg_len, self_sig_sk);
	
	//return
	TEE_MemMove(params[2].memref.buffer, sig_ct, sig_ct_len);
	TEE_MemMove(params[3].memref.buffer, &sig_ct_len, sizeof(size_t));
	// TEE_Free(sig_msg);
	TEE_Free(sig_ct);
	TEE_Free(self_sig_sk);
}

TEE_Result ta_sign_verify(uint32_t param_types, TEE_Param params[4])
{
	const uint32_t exp_param_types =
		TEE_PARAM_TYPES(TEE_PARAM_TYPE_MEMREF_INPUT,
			TEE_PARAM_TYPE_MEMREF_INPUT,
			TEE_PARAM_TYPE_MEMREF_OUTPUT,
			TEE_PARAM_TYPE_MEMREF_OUTPUT
			);
	TEE_Result res;
	
	if (param_types != exp_param_types)
		return TEE_ERROR_BAD_PARAMETERS;
	
	//sign_verify
	uint8_t oppo_sig_pk[CRYPTO_PUBLICKEYBYTES_SIG];
	uint8_t oppo_id[ID_LEN];
	uint8_t *auth_result;
	auth_result = TEE_Malloc(2, 0);
	if(!auth_result)
		return TEE_ERROR_OUT_OF_MEMORY;
	size_t auth_result_len=0;
	uint8_t *sig_msg;
	int ret=0;
	size_t sig_msg_len = 2 + CRYPTO_BYTES_SIG;
	
	sig_msg = TEE_Malloc(sig_msg_len, 0);
	if(!sig_msg)
		return TEE_ERROR_OUT_OF_MEMORY;
	
	TEE_MemMove(sig_msg, params[0].memref.buffer, sig_msg_len);
	TEE_MemMove(oppo_id, params[1].memref.buffer, ID_LEN);
	// return ret;

	ret = sig_data_search(&sig_node_head, oppo_id, oppo_sig_pk);
	if(ret) {
      printf("Found failed\n");
	  return 404;
    }

	ret = crypto_sign_open(auth_result, &auth_result_len, sig_msg, sig_msg_len, oppo_sig_pk);
	if(ret) {
      printf("Verification failed\n");
	  return 403;
    }

	//return
	TEE_MemMove(params[2].memref.buffer, auth_result, auth_result_len);
	TEE_MemMove(params[3].memref.buffer, &auth_result_len, sizeof(size_t));
	TEE_Free(sig_msg);
	TEE_Free(auth_result);
	// TEE_Free(oppo_sig_pk);
	// TEE_Free(&auth_result_len);
}

TEE_Result ta_context_sign_verify(uint32_t param_types, TEE_Param params[4])
{
	const uint32_t exp_param_types =
		TEE_PARAM_TYPES(TEE_PARAM_TYPE_MEMREF_INPUT,
			TEE_PARAM_TYPE_MEMREF_INPUT,
			TEE_PARAM_TYPE_MEMREF_OUTPUT,
			TEE_PARAM_TYPE_MEMREF_OUTPUT
			);
	TEE_Result res;
	
	if (param_types != exp_param_types)
		return TEE_ERROR_BAD_PARAMETERS;
	
	//sign_verify
	uint8_t oppo_sig_pk[CRYPTO_PUBLICKEYBYTES_SIG];
	uint8_t oppo_id[ID_LEN];
	uint8_t *auth_result;
	auth_result = TEE_Malloc(2, 0);
	if(!auth_result)
		return TEE_ERROR_OUT_OF_MEMORY;
	size_t auth_result_len=0;
	uint8_t *sig_msg;
	int ret=0;
	size_t sig_msg_len = 2 + CRYPTO_BYTES_SIG;
	
	sig_msg = TEE_Malloc(sig_msg_len, 0);
	if(!sig_msg)
		return TEE_ERROR_OUT_OF_MEMORY;
	
	TEE_MemMove(sig_msg, params[0].memref.buffer, sig_msg_len);
	TEE_MemMove(oppo_id, params[1].memref.buffer, ID_LEN);

	ret = sig_data_search(&sig_node_head, oppo_id, oppo_sig_pk);
	if(ret) {
      printf("Found failed\n");
	  return 404;
    }

	ret = crypto_sign_open(auth_result, &auth_result_len, sig_msg, sig_msg_len, oppo_sig_pk);
	if(ret) {
      printf("Verification failed\n");
	  return 403;
    }

	//return
	TEE_MemMove(params[2].memref.buffer, auth_result, auth_result_len);
	TEE_MemMove(params[3].memref.buffer, &auth_result_len, sizeof(size_t));
	TEE_Free(sig_msg);
	TEE_Free(auth_result);
	// TEE_Free(oppo_sig_pk);
	// TEE_Free(&auth_result_len);
}

TEE_Result TA_InvokeCommandEntryPoint(void __unused *sess_ctx,
			uint32_t __unused cmd_id,
			uint32_t __unused param_types,
			TEE_Param __unused params[4])
{
	switch (cmd_id)
	{
	case TA_SAVE_OPPOSIGPK:
		return ta_save_opposigpk(param_types, params);
	case TA_SIG_KEYPAIR:
		return ta_sig_keypair(param_types, params);
	case TA_KEM_KEYPAIR://Alice generate pk sk, sign the pk with Ksa
		return ta_kem_keypair(param_types, params);
	case TA_AUTHEN_ENCAPE://Bob verify the pk with Kva, and encape the K, and sign the ciphertext with Ksb
		return ta_authen_encape(param_types, params);
	case TA_AUTHEN_DECAPE://Alice verify the ciphertext with Kvb, and decape the k
		return ta_authen_decape(param_types, params);
	case TA_SIGN:
		return ta_sign(param_types, params);
	case TA_SIGN_VERIFY:
		return ta_sign_verify(param_types, params);
	case TA_CONTEXT_SIGN_VERIFY:
		return ta_context_sign_verify(param_types, params);
	default:
		break;
	}
	return TEE_ERROR_BAD_PARAMETERS;
}